#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

# !/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from tqdm import tqdm

valid_chrom = set(
    ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17", "18", "19", "20",
     "21", "22"])

client = MongoClient("mongodb://han:han2019@mongo-han.handb:27017")
collection = client["hanvcf"]["all_20210521"]

infile = open("variation_info_GRCh38.txt", "r")
outfile = open("variation_info_GRCh38_addhan.tsv", "w")
outfile.write(infile.readline().rstrip("\n") + "\tPggHan Sample Size\tPggHan Ref Allele\tPggHan Alt Allele\n")
outfile.flush()
for line in tqdm(infile):
    line = line.rstrip("\n")
    splitline = line.split("\t")
    chrom = splitline[0]
    pos = int(splitline[1])
    if chrom in valid_chrom:
        chrom = int(chrom)
        result_count = collection.count_documents({"chrom": chrom, "pos": pos})
        if result_count == 0:
            outfile.write(line + "\tnull\tnull\tnull\n")
        elif result_count == 1:
            result = collection.find_one({"chrom": chrom, "pos": pos})
            ref = result["ref"]
            alt = result["alt"]
            individual_count = 0
            ref_count = 0
            alt_count = 0
            for item in result["region"]:
                individual_count += item["individual_count"]
                ref_count += item["rr"] * 2
                ref_count += item["ra"]
                alt_count += item["ra"]
                alt_count += item["aa"] * 2
            individual_count = individual_count * 2
            alt_freq = float(alt_count) / float(individual_count)
            ref_freq = float(ref_count) / float(individual_count)
            ref_str = "%s=%f" % (ref, ref_freq)
            alt_str = "%s=%f" % (alt, alt_freq)

            outfile.write(line + "\t%d\t%s\t%s\n" % (individual_count, ref_str, alt_str))
        else:
            results = collection.find({"chrom": chrom, "pos": pos})
            individual_count = 0
            ref_count = 0
            ref = ""
            alt_list = {}
            for result in results:
                ref = result["ref"]
                alt = result["alt"]
                alt_count = 0
                for item in result["region"]:
                    individual_count += item["individual_count"]
                    ref_count += item["rr"] * 2
                    ref_count += item["ra"]
                    alt_count += item["ra"]
                    alt_count += item["aa"] * 2
                alt_list[alt] = alt_count
            individual_count = individual_count * 2
            ref_freq = float(ref_count) / float(individual_count)
            alt_str_list = []
            for alt in alt_list:
                alt_freq = float(alt_list[alt]) / float(individual_count)
                alt_str_list.append("%s=%f" % (alt, alt_freq))
            ref_str = "%s=%f" % (ref, ref_freq)
            alt_str = "|".join(alt_str_list)
            outfile.write(line + "\t%d\t%s\t%s\n" % (individual_count, ref_str, alt_str))
    else:
        outfile.write(line + "\tnull\tnull\tnull\n")
    outfile.flush()
infile.close()
outfile.close()
